import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import math
import copy  # 用于SimpleDataset深拷贝
from torch.utils.data import Dataset, DataLoader
import pathlib
import os

# -------------------------- 1. 基础配置（设备、随机种子） --------------------------
# 设置随机种子
seed = 42
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

# 设备配置（用户原代码中的Device变量，统一为device）
Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using Device: {Device}")

save_model_folder = pathlib.Path(__file__).parent.joinpath('models')
save_model_folder.mkdir(parents=True, exist_ok=True)
# -------------------------- 2. 用户指定的SimpleDataset（固定长度数据） --------------------------
class SimpleDataset(Dataset):
    def __init__(self, num):
        prefix = torch.tensor([0])  # 前缀：起始符号0
        suffix = torch.tensor([9])  # 后缀：结束符号9
        # 输入数据：num个(5,)的tensor，值范围1-8
        self.data = [torch.randint(1, 9, size=(5,), dtype=torch.long) for _ in range(num)]
        # 标签：前缀+数据+后缀 → 每个label是(7,)的tensor（0 + 5个数据 + 9）
        self.labels = [torch.cat([prefix, _, suffix], dim=0) for _ in copy.deepcopy(self.data)]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]  # data: (5,), label: (7,)


# 用户指定的创建DataLoader函数（固定长度无需填充，简化collate）
def create_data(batch_size=5, num=100):
    dataset = SimpleDataset(num)
    # 固定长度数据，无需自定义collate_fn，用默认即可（直接堆叠）
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader


# -------------------------- 3. Transformer核心模块（保留原逻辑，适配固定长度） --------------------------
# 3.1 位置编码（支持最大序列长度10，覆盖data(5)、label(7)）
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=10, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 生成位置编码：(max_len, d_model)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))  # (d_model//2,)

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位用正弦
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位用余弦

        pe = pe.unsqueeze(1)  # (max_len, 1, d_model) → 适配(batch_size)维度
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: (seq_len, batch_size, d_model) → 加上对应长度的位置编码
        x = x + self.pe[:x.size(0), :, :]
        return self.dropout(x)


# 3.2 编码器层（复用原多头注意力逻辑）
class EncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # src: (seq_len, batch_size, d_model)
        src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)  # 残差连接
        src = self.norm1(src)  # 层归一化

        # 前馈网络
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)  # 残差连接
        src = self.norm2(src)  # 层归一化
        return src


# 3.3 解码器层（复用原逻辑，生成下三角掩码防止看未来）
class DecoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)  # 解码器自注意力
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)  # 编解码交叉注意力

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = F.relu

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                tgt_key_padding_mask=None, memory_key_padding_mask=None):
        # tgt: (tgt_seq_len, batch_size, d_model)；memory: (src_seq_len, batch_size, d_model)
        # 1. 解码器自注意力（带掩码，防止看未来token）
        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        # 2. 编解码交叉注意力（用编码器输出memory）
        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)

        # 3. 前馈网络
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt


# 3.4 完整Transformer模型（适配用户train()中的模型参数）
class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, seq_len, d_model=512,
                 dim_feedforward=2048, num_layers=6, nhead=8, dropout=0.1):
        """
        适配用户train()中的模型初始化：Transformer(11,11,10,512,4*512,6,8)
        - src_vocab_size: 源词汇表大小（0-9 → 10，用户传11不影响，预留1位）
        - tgt_vocab_size: 目标词汇表大小（同src）
        - seq_len: 最大序列长度（用户传10，足够覆盖5/6/7长度）
        - d_model: 嵌入维度（512）
        - dim_feedforward: 前馈网络维度（4*512=2048）
        - num_layers: 编解码器层数（6）
        - nhead: 多头注意力头数（8）
        """
        super(Transformer, self).__init__()
        # 嵌入层（源和目标共用同一嵌入层，也可分开，这里简化）
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.d_model = d_model

        # 位置编码
        self.pos_encoder = PositionalEncoding(d_model, max_len=seq_len, dropout=dropout)

        # 编码器（num_layers层EncoderLayer）
        encoder_layer = EncoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # 解码器（num_layers层DecoderLayer）
        decoder_layer = DecoderLayer(d_model, nhead, dim_feedforward, dropout)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # 输出层（映射到目标词汇表）
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)

        # 初始化权重
        self.init_weights()

    def init_weights(self):
        # 嵌入层和输出层权重初始化
        initrange = 0.1
        self.src_embedding.weight.data.uniform_(-initrange, initrange)
        self.tgt_embedding.weight.data.uniform_(-initrange, initrange)
        self.fc_out.bias.data.zero_()
        self.fc_out.weight.data.uniform_(-initrange, initrange)

    def generate_tgt_mask(self, tgt_seq_len):
        """生成解码器自注意力掩码：下三角矩阵（防止看未来token）"""
        mask = nn.Transformer.generate_square_subsequent_mask(tgt_seq_len).to(Device)
        return mask  # (tgt_seq_len, tgt_seq_len)

    def forward(self, src, tgt):
        """
        适配用户train()的输入格式：
        - src: 源输入（batch_size, src_seq_len）→ 对应SimpleDataset的data: (batch_size, 5)
        - tgt: 目标输入（batch_size, tgt_seq_len）→ 对应input_label: (batch_size, 6)
        返回：output (batch_size, tgt_seq_len, tgt_vocab_size)
        """
        src_seq_len = src.size(1)  # 5
        tgt_seq_len = tgt.size(1)  # 6
        batch_size = src.size(0)  # 批次大小

        # 1. 嵌入层 + 位置编码（需转置为(seq_len, batch_size, d_model)，适配PyTorch Transformer）
        src_emb = self.src_embedding(src).transpose(0, 1) * math.sqrt(self.d_model)  # (5, batch_size, 512)
        src_emb = self.pos_encoder(src_emb)

        tgt_emb = self.tgt_embedding(tgt).transpose(0, 1) * math.sqrt(self.d_model)  # (6, batch_size, 512)
        tgt_emb = self.pos_encoder(tgt_emb)

        # 2. 生成解码器掩码（下三角掩码，防止看未来token）
        tgt_mask = self.generate_tgt_mask(tgt_seq_len)  # (6, 6)

        # 3. 编码器前向传播（无pad，src_key_padding_mask=None）
        memory = self.encoder(src_emb, src_key_padding_mask=None)  # (5, batch_size, 512)

        # 4. 解码器前向传播（无pad，tgt_key_padding_mask=None）
        tgt_out = self.decoder(
            tgt_emb, memory,
            tgt_mask=tgt_mask,
            tgt_key_padding_mask=None,
            memory_key_padding_mask=None
        )  # (6, batch_size, 512)

        # 5. 输出层（转置回(batch_size, tgt_seq_len, tgt_vocab_size)，适配用户损失计算）
        output = self.fc_out(tgt_out).transpose(0, 1)  # (batch_size, 6, tgt_vocab_size)
        return output


# -------------------------- 4. 用户指定的train()函数（完整适配） --------------------------
def train():
    
    model_path = save_model_folder.joinpath(f'model.pth').resolve().__str__()
    
    # 1. 数据集配置（用户指定：batch_size=32，样本数=1000）
    datas = create_data(batch_size=32, num=1000)

    # 2. 模型初始化（完全匹配用户原代码参数：Transformer(11,11,10,512,4*512,6,8)）
    
    if os.path.exists(model_path):  # 模型已存在
        print('模型存在了，直接加载')
        model = torch.load(model_path)
    else:
        model = Transformer(
            src_vocab_size=11,  # 源词汇表大小（用户传11，覆盖0-9+预留1位）
            tgt_vocab_size=11,  # 目标词汇表大小（同src）
            seq_len=10,  # 最大序列长度（用户传10）
            d_model=512,  # 嵌入维度（512）
            dim_feedforward=4 * 512,  # 前馈网络维度（2048）
            num_layers=6,  # 编解码器层数（6）
            nhead=8  # 多头注意力头数（8）
        ).to(Device)

    # 3. 训练配置（用户指定：SGD优化器、lr=1e-1、无学习率调度）
    intervals = 30  # 训练轮次（用户指定）
    criterion = nn.CrossEntropyLoss()  # 无pad，无需ignore_index
    opt = optim.SGD(model.parameters(), lr=1e-1)

    # 4. 训练循环（完全保留用户结构）
    for interval in range(intervals):
        model.train()
        loss_all = 0.0

        for data, label in datas:
            # 数据移动到设备（data: (32,5), label: (32,7)）
            data = data.to(Device)
            label = label.to(Device)

            # 目标输入/输出拆分（用户原逻辑：label[:,:-1]为解码器输入，label[:,1:]为损失计算目标）
            input_label = label[:, :-1]  # (32,6) → 解码器输入（去掉最后一位9）
            output_label = label[:, 1:]  # (32,6) → 损失计算目标（去掉第一位0）

            # 梯度清零 + 前向传播
            opt.zero_grad()
            output = model(data, input_label)  # output: (32,6,11)

            # 计算损失（reshape为(batch_size*tgt_seq_len, vocab_size) vs (batch_size*tgt_seq_len,)）
            loss = criterion(
                output.reshape(-1, output.size(-1)),  # (32*6,11) = (192,11)
                output_label.reshape(-1)  # (32*6,) = (192,)
            )
            loss_all += loss.item()

            # 反向传播 + 梯度裁剪 + 优化器更新
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 防止梯度爆炸
            opt.step()

        # 计算平均损失并打印（保留用户格式）
        avg_loss = loss_all / len(datas)
        print(f'第 {interval + 1} 轮, 平均损失: {avg_loss:.6f}, 学习率: {opt.param_groups[0]["lr"]:.8f}')
        torch.save(model, model_path)
        
    with torch.no_grad():
        # 测试数据
        test_in_datas = [1, 2, 3, 4, 5]
        test_in_data = torch.Tensor(test_in_datas).long().view(1, -1)

        # label = [1,2,3,4,5]
        label = [0]
        label = torch.Tensor(label).long().view(1, -1)

        while 1:
            test_in_data = test_in_data.to(Device)
            label = label.to(Device)

            mask = torch.triu(torch.ones(len(label), len(label)), diagonal=1) == 0
            # mask = mask.view(1,1,1,1)
            mask_self = mask.to(Device)

            output = model(test_in_data, label)
            log_softmax = F.log_softmax(output, dim=-1)
            # print(log_softmax)
            max_value = log_softmax.argmax(dim=-1)
            # label = max_value[-1]
            label = torch.cat([label, max_value[:, -1].view(1, -1)], dim=1)
            print(label)
            if max_value[:, -1].item() == 9:  # 9为结束符
                break




# -------------------------- 5. 启动训练 --------------------------
if __name__ == "__main__":
    train()